import os
import sys

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
import torch.nn as nn
import torch.optim as optim
from os.path import dirname, abspath

ppfolder = dirname(dirname(abspath(__file__)))
if (ppfolder not in sys.path):
    sys.path.insert(0, ppfolder)

from tools import add_project_root_2_sys_path

add_project_root_2_sys_path()

from data import get_dataloader_train
from models import MnistNet
from data import FashionMNIST_CLASS_LIST

from data import TBV

if __name__ == '__main__':
    FashionMnistFolder = os.path.join(r"D:\deepdotdev\deepdot-vision\datasets\FashionMNIST\raw")
    dataloader_train = get_dataloader_train(FashionMnistFolder, image_size=(28, 28), batch_size=1010)

    net = MnistNet()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
    classes = FashionMNIST_CLASS_LIST

    tb_vis = TBV(r"D:\deepdotdev\deepdot-vision\runs\00")

    dataiter = iter(dataloader_train)
    images, labels = dataiter.next()

    tb_vis.add_image_batch(images, 0, "FashionImages")
    tb_vis.add_graph_image(net, images)
    for epoch in range(100):
        loss_mean = 0.
        loss_sum_epoch = 0.
        loss_count = 0
        correct = 0.
        total = 0.

        net.train()

        for i, data in enumerate(dataloader_train):
            # forward
            inputs, labels = data
            outputs = net(inputs)

            # Compute loss
            optimizer.zero_grad()
            loss = criterion(outputs, labels)

            loss_sum_epoch += loss
            loss_count += 1

            # backward
            loss.backward()

            # updata weights
            optimizer.step()

        loss_mean = loss_sum_epoch / loss_count
        tb_vis.add_loss_step(loss_mean, epoch, "LossValues")
        print("current epoch is {} with loss {}".format(epoch, loss_mean))
